function M = rhino_icp(data1,data2,Ninits,ax, hs_nasion, pol_nasion)
% Run Iterative Closest Point with multiple initialisations
% M = rhino_icp(data1,data2,Ninits,do_plots
%
% REQUIRED INPUTS:
%
% data1      - [3 x N] locations of the set of static points 
%              (e.g. MRI scalp surface)
% data2      - [3 x N] locations of the set of moving points 
%              (e.g. Polhemus headshape points)
%
% OPTIONAL INPUTS:
%
% Ninits     - Number of random initialisations to perform (default 10)
% ax         - Axes handle or [] to plot iterations
% hs_nasion  - nasion in MRI surface
% pol_nasion - nasion in pol points    -- these last two constrain fit to
%                                         match these points within 5mm. 
%
%
% OUTPUTS:
%
% M          - [4 x 4] rigid transformation matrix mapping data2 to data1
%
% Based on ICP.m by Martin Kjer and Jakob Wilm, Technical University of 
% Denmark
%
% Adam Baker 2014

if nargin < 3
    Ninits = 10;
end
if nargin < 4
    do_plots = 0;
else
    do_plots = 1;
    if isempty(ax)
        ax = axes;
    end
end
  
% identify if we want to fix nasion position. 
% We use this if we have legacy Neuromag data, without a nose, but where
% the nasion is likely known to high accuracy
if nargin > 4 && exist('hs_nasion', 'var') && exist('pol_nasion', 'var') ...
              && ~(isempty(hs_nasion)      || isempty(pol_nasion)),
    useICPconstrained = true;
else
    useICPconstrained = false;
end%if

            
err_old = Inf;
err     = zeros(1,Ninits);

Mr      = eye(4);
data2r  = data2;

for init = 1:Ninits,
    if useICPconstrained,
        % constrain search space to penalise movement of nasion outside of
        % 5mm sphere
        bound = 5; %mm
        [TR,TT,e] = icp_constrained(data1,data2r,50, ...
                                    'Minimize','lmaPointConstrained',...
                                    'q_constrainedPoint', hs_nasion, 'p_constrainedPoint', pol_nasion, 'pointBound', bound,...
                                    'Matching','kDtree','ReturnAll',true);
    else
        % unconstrained icp search space
        [TR,TT,e] = icp_constrained(data1,data2r,50,'Minimize','plane','Matching',...
                                   'kDtree','ReturnAll',true);
    end   
    
    % compare error to previous best fit
    [~,emin] = min(e);
    e  = e(1:emin);
    TR = TR(:,:,1:emin);
    TT = TT(:,:,1:emin);
    
    Mi = make_affine_transformation(TR(:,:,end), TT(:,:,end));

    err(init) = e(end);
 
    if err(init) < err_old
        err_old = err(init);
        M       = Mi*Mr;
    end
        
    % update plots showing quality of fit
    if do_plots
        hs = plot(ax,[]);
        xlim = get(ax,'xlim'); ylim = get(ax,'ylim'); zlim = get(ax,'zlim');
        for i = 1:length(e)
            delete(hs);
            Mi = [TR(:,:,i) TT(:,:,i);0 0 0 1];
            Mi = Mi*Mr;
            dataM = spm_eeg_inv_transform_points(Mi,data2')';
            hold(ax,'on')
            hs = plot3(ax,dataM(1,:),dataM(2,:),dataM(3,:),'.b');
            set(ax,'xlim',xlim); set(ax,'ylim',ylim); set(ax,'zlim',zlim);
            title(ax,num2str(e(i)))
            drawnow
        end
        delete(hs)
    end
    
    % Give the registration a kick...
    if useICPconstrained,
        Mr = constrained_initialisation(pol_nasion);
    else
        Mr = unconstrained_initialisation;
    end%if
        
    data2r  = transform_points(Mr, data2);
   
end% loop over initialisations
end%rhino_icp

function Mr = unconstrained_initialisation
% 15 degrees rotation about each axis, 10 mm translation
a = (rand-0.5)*pi/6; 
b = (rand-0.5)*pi/6; 
c = (rand-0.5)*pi/6;

Rx = [1          0     0   ;   0    cos(a) -sin(a);    0    sin(a) cos(a)];
Ry = [cos(b)     0   sin(b);   0      1       0   ; -sin(b)   0    cos(b)];
Rz = [cos(c) -sin(c)   0   ; sin(c) cos(c)    0   ;    0      0       1  ];

T  = 10*[(rand-0.5) (rand-0.5) (rand-0.5)];

Mr = make_affine_transformation(Rx*Ry*Rz, T(:));
end%unconstrained_initialisation

function Mr = constrained_initialisation(fixPoint)
% performs a constrained initialisation which does not move FIXPOINT. 
% corresponds to a rotation about the fixPoint - reduced DoF

% get a random rotation matrix without translation
M = unconstrained_initialisation;
R = M(1:3,1:3);

% restore fixed point using a translation
T = (eye(3) - R) * fixPoint(:);

Mr = make_affine_transformation(R, T);

% check that we've not moved the rotation 'origin'
assert(norm(fixPoint - transform_points(Mr, fixPoint)) < eps(max(fixPoint)), ...
       [mfilename ':ConstrainedInitialisationFailed'],                       ...
       'Constrained initialisation failed to keep fixed point. \n');
end%constrained_initialisation

function newPoints = transform_points(M, oldPoints)
% applies affine transformation to OLDPOINTS, 3xN

out       = M * [oldPoints; ones(1, size(oldPoints, 2))];
newPoints = out(1:3,:);
end%transform_points

function M = make_affine_transformation(R, T)
% makes affine transformation matrix from rotation matrix and
% transformation matrix
M = [R T; 0 0 0 1];
end%make_affine_transformation
